module fourier_space_diagnostic
use mod_2dflu
contains
!-------------------------------------
subroutine calc_fspace_quant(idiag,iflag)
use omplib
implicit none
integer :: i1,i2,k1,k2,ireal,iimag,ksqr,mshl
integer :: idiag,iflag
integer :: nthreads,id
double precision :: rk2,kx,ky,rk,one_by_nsqr,energy_norm
double precision :: psi_real,psi_imag
double complex :: oo,uu_x,uu_y,psi_complex
!! get psi and velocity from the omega
factor = 2.0d0*pi/length
one_by_nsqr = 1.0d0/dfloat(n1*n2)
energy_norm = one_by_nsqr*one_by_nsqr
!
!$OMP PARALLEL DO PRIVATE(k2,i1,k1,ireal,iimag,ksqr,rk2,kx,ky,psi_real,psi_imag,oo,uu_x,uu_y,psi_complex)  SHARED(omega,psi,ukx,uky)
do i2 = 1,n2
  k2 = (i2-1) - n2*(i2/(n1hf+1))	
  do i1 =1,n1hf
    k1 = i1 -1
    ksqr = k1*k1 + k2*k2 
    rk2 = factor*factor*dfloat(ksqr)
    ireal = 2*i1 -1
    iimag = 2*i1
    kx = factor*dfloat(k1)
    ky = factor*dfloat(k2)
    if((ksqr.gt.kasqr).or.(ksqr.eq.0))then
      omega(ireal,i2) = 0.0d0;omega(iimag,i2) = 0.0d0
      psi(ireal,i2) = 0.0d0;psi(iimag,i2) = 0.0d0
      ukx(ireal,i2) = 0.0d0;ukx(iimag,i2) = 0.0d0
      uky(ireal,i2) = 0.0d0;uky(iimag,i2) = 0.0d0
      if((idiag.eq.1).and.(iflag.eq.1)) then
        if(lcalc_Gradu) then
          dx_ux(ireal,i2) = 0.0d0; dx_ux(iimag,i2) = 0.0d0
          dy_ux(ireal,i2) = 0.0d0; dy_ux(iimag,i2) = 0.0d0
          dx_uy(ireal,i2) = 0.0d0; dx_uy(iimag,i2) = 0.0d0
        endif
      endif
    else
      oo=dcmplx(omega(ireal,i2),omega(iimag,i2))
!
! \lap \psi = -\omega, -k^2 \hat{\psi} = -\hat{\omega},\hat{\psi} = \hat{\omega}/k^2
!
      psi_complex=oo/rk2
!
! u_x = \partial_y \psi, \hat{u}_x = i ky \hat{psi}
!
      uu_x = zi*ky*psi_complex
!
! u_y = -\partial_x \psi, \hat{u}_x = -i kx \hat{psi}
!
      uu_y = -zi*kx*psi_complex
!
      psi(ireal,i2) = real(psi_complex);psi(iimag,i2) = aimag(psi_complex)
      ukx(ireal,i2) = real(uu_x);       ukx(iimag,i2) = aimag(uu_x) 
      uky(ireal,i2) = real(uu_y);       uky(iimag,i2) = aimag(uu_y)
      if((idiag.eq.1).and.(iflag.eq.1)) then
!
!Energy and vorticity should be calculated here, but that would involve some 
!reduction in shared memory parallisation. So we have not coded that here yet. 
!
        if(lcalc_Gradu) then
          dx_ux(ireal,i2) = -kx*ky*real(psi_complex); dx_ux(iimag,i2) = -kx*ky*aimag(psi_complex)
          dy_ux(ireal,i2) = -ky*ky*real(psi_complex); dy_ux(iimag,i2) = -ky*ky*aimag(psi_complex)
          dx_uy(ireal,i2) = kx*kx*real(psi_complex); dx_uy(iimag,i2) = kx*kx*aimag(psi_complex)
        endif
      endif
    endif
  enddo
enddo
!$OMP END PARALLEL DO
!! --------------------- 
end subroutine calc_fspace_quant
!! --------------------- 
end module fourier_space_diagnostic
